package com.gitee.wsl.concurrent.api.cancel

import com.gitee.wsl.platform.exception.handleReaktiveError
import kotlinx.atomicfu.AtomicRef
import kotlinx.atomicfu.atomic
import kotlin.jvm.JvmOverloads
import kotlin.jvm.JvmRecord

/**
 * INTERNAL API.
 *
 *
 * **WARN:** Internal apis are subject to change or removal without
 * any notice. When code depends on internal APIs, it is subject to breakage
 * between minor version updates.
 */
class MutableCancellable @JvmOverloads constructor(
    initialRef: Cancellable = Cancellable.EMPTY
) : Cancellable {
    private val ref: AtomicRef<State> = atomic<State>(State.Active(initialRef, 0, null))

    override fun cancel() {
        var state:State? = ref.getAndSet(State.Closed)
        while (state is State.Active) {
            try {
                state.token?.cancel()
            } catch (e: Exception) {
                e.handleReaktiveError()
            }
            state = state.rest
        }
    }

    fun newCancellableRef(): CancellableForwardRef {
        when (val current = ref.value) {
            is State.Closed -> {
                return CancellableForwardRef { cancellable -> cancellable.cancel() }
            }

            is State.Active -> {
                return object : CancellableForwardRef {
                    override fun set(cancellable: Cancellable) {
                        registerOrdered(
                            current.order,
                            cancellable,
                            current
                        )
                    }

                    override fun cancel() {
                        unregister(current.order)
                    }
                }
            }

            else -> {
                throw IllegalStateException("Invalid state: $current")
            }
        }
    }

    fun register(token: Cancellable): Cancellable? {
        while (true) {
            val current = ref.value
            if (current is State.Active) {
                val newOrder = current.order + 1
                val update = State.Active(token, newOrder, current)
                if (ref.compareAndSet(current, update)) {
                    return Cancellable { unregister(newOrder) }
                }
            } else if (current is State.Closed) {
                token.cancel()
                return null
            } else {
                throw IllegalStateException("Invalid state: $current")
            }
        }
    }

    private fun unregister(order: Long) {
        while (true) {
            val current = ref.value
            if (current is State.Active) {
                var cursor: State.Active? = current
                var acc: State.Active? = null
                while (cursor != null) {
                    if (cursor.order != order) {
                        acc = State.Active(cursor.token, cursor.order, acc)
                    }
                    cursor = cursor.rest
                }
                // Reversing
                var update: State.Active? = null
                while (acc != null) {
                    update = State.Active(acc.token, acc.order, update)
                    acc = acc.rest
                }
                if (update == null) {
                    update = State.Active(Cancellable.Companion.EMPTY, 0, null)
                }
                if (ref.compareAndSet(current, update)) {
                    return
                }
            } else if (current is State.Closed) {
                return
            } else {
                throw IllegalStateException("Invalid state: $current")
            }
        }
    }

    private fun registerOrdered(
        order: Long,
        newToken: Cancellable,
        current: State?,
    ) {
        var current = current
        while (true) {
            if (current is State.Active) {
                // Double-check ordering
                if (current.order != order) {
                    return
                }
                // Try to update
                val update = State.Active(newToken, order + 1, null)
                if (ref.compareAndSet(current, update)) {
                    return
                }
                // Retry
                current = ref.value
            } else if (current is State.Closed) {
                newToken.cancel()
                return
            } else {
                throw IllegalStateException("Invalid state: $current")
            }
        }
    }

    internal interface State {
        @JvmRecord
        data class Active(
            val token: Cancellable?,
            val order: Long,
            val rest: Active?,
        ) : State

        object Closed : State
    }
}
